import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

# 1. 数据准备
# 示例数据
data = pd.read_csv("../demo/2024-2.csv", encoding='GBK')
df = pd.DataFrame(data)

# 将日期列转换为日期类型
df['date_id'] = pd.to_datetime(df['date_id'], format='%Y%m%d')

# 2. 数据探索
# 按 SKU 分组绘制销量趋势
for sku, group in df.groupby('sku'):
    plt.plot(group['date_id'], group['sales'], label=f'SKU {sku}')

plt.xlabel('date_id')
plt.ylabel('Sales')
plt.title('Sales Trend by SKU')
plt.legend()
plt.show()

# 3. 特征工程
# 添加时间特征
df['year'] = df['date_id'].dt.year
df['month'] = df['date_id'].dt.month
df['day'] = df['date_id'].dt.day
df['day_of_week'] = df['date_id'].dt.dayofweek

# 添加滞后特征
df['lag_1'] = df.groupby('sku')['sales'].shift(1)  # 前一天的销量
df['lag_2'] = df.groupby('sku')['sales'].shift(2)  # 前两天的销量

# 添加滑动平均
df['rolling_mean_3'] = df.groupby('sku')['sales'].rolling(window=3).mean().reset_index(level=0, drop=True)

# 删除缺失值（由于滞后特征引入）
df = df.dropna()

# 4. 模型训练与评估
# 按 SKU 分组训练和预测
results = {}
for sku, group in df.groupby('sku'):
    X = group.drop(columns=['date_id', 'sku', 'sales'])
    y = group['sales']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

    # 训练 LightGBM 模型
    model = lgb.LGBMRegressor()
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # # 计算 RMSE
    # rmse = mean_squared_error(y_test, y_pred, squared=False)
    # results[sku] = rmse
    #
    # print(f'SKU {sku} RMSE: {rmse}')

    # 可视化预测结果
    plt.plot(group['date_id'], y, label='Actual')
    plt.plot(group.loc[X_test.index, 'date_id'], y_pred, label='Predicted', linestyle='--')
    plt.title(f'SKU {sku} Sales Prediction')
    plt.legend()
    plt.show()

# 输出结果
print("RMSE by SKU:", results)

# 5. 未来销量预测
# 构建未来日期
future_dates = pd.date_range(start=df['date_id'].max() + pd.Timedelta(days=1), periods=7)  # 预测未来7天
future_data = pd.DataFrame({'date_id': future_dates})

# 添加特征
future_data['year'] = future_data['date_id'].dt.year
future_data['month'] = future_data['date_id'].dt.month
future_data['day'] = future_data['date_id'].dt.day
future_data['day_of_week'] = future_data['date_id'].dt.dayofweek

# 预测每个 SKU 的未来销量
for sku, group in df.groupby('sku'):
    # 获取最后一个有效数据点
    last_row = group.iloc[-1]

    # 添加滞后特征和滑动平均
    future_data['lag_1'] = last_row['sales']
    future_data['lag_2'] = last_row['lag_1']
    future_data['rolling_mean_3'] = last_row['rolling_mean_3']

    # 预测
    future_data[f'sales_pred_{sku}'] = model.predict(future_data.drop(columns=['date_id']))

# 输出未来预测结果
print(future_data[['date_id'] + [col for col in future_data.columns if 'sales_pred' in col]])